import pandas as pd
import numpy as np



class BiasSVD:

    def __init__(self, learning_rate=0.01, reg_rate=0.1, n_epochs=100, n_factors=15, verbose=True):
        """
        FunkSVD算法的初始化方法
        :param learning_rate: 学习率
        :param reg_rate: 正则化系数
        :param n_epochs: 迭代次数
        :param n_factors: 隐特征向量维度
        :param verbose: 是否打印训练过程log
        """
        self.learning_rate = learning_rate
        self.reg_rate = reg_rate
        self.n_epochs = n_epochs
        self.n_factors = n_factors
        self.verbose = verbose

    def fit(self, X, val_data=None):
        """
        FunkSVD算法的训练方法
        :param X: 用户-物品评分矩阵
        :param val_data: 验证数据，用于退出训练
        :return: 用户特征矩阵P和物品特征矩阵Q的乘积
        """
        # 初始化变量
        n_users, n_items = X.shape
        self.P = np.random.normal(size=(n_users, self.n_factors))
        self.Q = np.random.normal(size=(n_items, self.n_factors))
        self.bu=np.zeros(X.shape)
        self.bi = np.zeros(X.shape)
        self.mean=X.mean().mean()
        # 训练模型
        for epoch in range(self.n_epochs):
            for i in range(n_users):
                for j in range(n_items):
                    if X.iloc[i,j] > 0:
                        error = X.iloc[i,j] - np.dot(self.P[i, :], self.Q[j, :].T)-self.mean-self.bi[i][j]-self.bu[i][j]
                        self.P[i, :] += self.learning_rate * (error * self.Q[j, :] - self.reg_rate * self.P[i, :])
                        self.Q[j, :] += self.learning_rate * (error * self.P[i, :] - self.reg_rate * self.Q[j, :])
                        self.bu[i][j]+= self.learning_rate * (error - self.reg_rate * self.bu[i][j])
                        self.bi[i][j] += self.learning_rate * (error - self.reg_rate * self.bi[i][j])
            # 计算验证集误差
            if val_data is not None:
                y_true, y_pred = self.evaluate(val_data)
                val_loss = np.sqrt(np.mean(np.power(y_true - y_pred, 2)))
                if self.verbose:
                    print('epoch %d, val_loss=%.4f' % (epoch + 1, val_loss))

        return np.dot(self.P, self.Q.T)

    def predict(self, u, i):
        """
        预测用户对物品的评分
        :param u: 用户ID
        :param i: 物品ID
        :return: 预测评分得分
        """
        return np.dot(self.P[u, :], self.Q[i, :].T)

    def evaluate(self, test_data):
        """
        评估算法效果
        :param test_data: 测试数据
        :return: 真实值和预测值
        """

        y_true = test_data['rating']
        y_pred = [self.predict(row['userId'], row['movieId']) for i, row in test_data.iterrows()]
        return y_true, y_pred

ratings = pd.read_csv('train.txt', encoding='GBK', sep=',')

# 将评分矩阵转换为二维矩阵形式
train = ratings.pivot_table(values='rating', index='userId', columns='movieId').fillna(0)
test=pd.read_csv('test.txt', encoding='GBK', sep=',')


# 初始化 FunkSVD 模型
Bias_svd = BiasSVD(n_epochs=50, n_factors=30, verbose=True)


# 训练模型
Bias_svd.fit(train, val_data=test)



print(np.dot(Bias_svd.P, Bias_svd.Q.T))
# 电影推荐
user_id = 1
user_ratings = ratings[ratings['userId'] == user_id]
user_unseen_movies = ratings[~ratings['movieId'].isin(user_ratings['movieId'])]['movieId'].unique()
recommendations = []

for movie_id in user_unseen_movies:
    rating = Bias_svd.predict(user_id - 1, movie_id - 1)
    recommendations.append((movie_id, rating))

# 推荐电影排序并输出
recommendations = sorted(recommendations, key=lambda x: x[1], reverse=True)[:10]
recommended_movie_ids = [rec[0] for rec in recommendations]
recommended_movies = ratings[ratings['movieId'].isin(recommended_movie_ids)]['title'].unique()

print('为用户推荐的电影有：')
for movie in recommended_movies:
    print(movie)